package com.srmt.common.ratelimit.impl;

import com.srmt.common.ratelimit.RateLimiter;
import com.srmt.common.ratelimit.config.RateLimitConfig;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Mock限流器实现，用于测试场景
 *
 * 特点：
 * - 可配置是否允许请求通过
 * - 记录调用次数统计
 * - 支持模拟限流场景
 *
 * @author SRMT Development Team
 * @since 2.0.0
 */
public class MockRateLimiter implements RateLimiter {

    private final RateLimitConfig config;
    private final Map<String, Integer> qpsMap = new ConcurrentHashMap<>();
    private final Map<String, AtomicInteger> acquireCountMap = new ConcurrentHashMap<>();

    // 测试控制参数
    private volatile boolean allowAcquire = true;  // 是否允许获取令牌
    private volatile int maxAllowedAcquires = Integer.MAX_VALUE;  // 最大允许获取次数
    private volatile long simulatedDelayMs = 0;  // 模拟延迟时间

    public MockRateLimiter(RateLimitConfig config) {
        this.config = config;
    }

    /**
     * 设置是否允许获取令牌（用于测试失败场景）
     */
    public void setAllowAcquire(boolean allowAcquire) {
        this.allowAcquire = allowAcquire;
    }

    /**
     * 设置最大允许获取次数（用于测试配额用尽场景）
     */
    public void setMaxAllowedAcquires(int maxAllowedAcquires) {
        this.maxAllowedAcquires = maxAllowedAcquires;
    }

    /**
     * 设置模拟延迟时间（用于测试超时场景）
     */
    public void setSimulatedDelayMs(long delayMs) {
        this.simulatedDelayMs = delayMs;
    }

    /**
     * 获取API的总调用次数
     */
    public int getAcquireCount(String api) {
        AtomicInteger counter = acquireCountMap.get(api);
        return counter != null ? counter.get() : 0;
    }

    /**
     * 重置调用计数
     */
    public void resetAcquireCount(String api) {
        if (api == null) {
            acquireCountMap.clear();
        } else {
            acquireCountMap.remove(api);
        }
    }

    @Override
    public boolean tryAcquire(String api, long timeout) {
        // 记录调用次数
        AtomicInteger counter = acquireCountMap.computeIfAbsent(api, k -> new AtomicInteger(0));
        int currentCount = counter.incrementAndGet();

        // 模拟延迟
        if (simulatedDelayMs > 0) {
            try {
                Thread.sleep(Math.min(simulatedDelayMs, timeout));
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return false;
            }
        }

        // 检查是否超过最大允许次数
        if (currentCount > maxAllowedAcquires) {
            return false;
        }

        // 返回配置的结果
        return allowAcquire;
    }

    @Override
    public void acquire(String api) throws InterruptedException {
        // 阻塞等待，直到获取成功
        while (!tryAcquire(api, 100)) {
            Thread.sleep(50);
        }
    }

    @Override
    public void setQps(String api, int qps) {
        qpsMap.put(api, qps);
    }

    @Override
    public int getQps(String api) {
        return qpsMap.getOrDefault(api, config.getDefaultQps());
    }

    @Override
    public int getAvailablePermits(String api) {
        int acquireCount = getAcquireCount(api);
        int maxPermits = maxAllowedAcquires;
        return Math.max(0, maxPermits - acquireCount);
    }

    @Override
    public void reset(String api) {
        if (api == null) {
            qpsMap.clear();
            acquireCountMap.clear();
        } else {
            qpsMap.remove(api);
            acquireCountMap.remove(api);
        }
    }

    @Override
    public void shutdown() {
        qpsMap.clear();
        acquireCountMap.clear();
    }
}
